"""
    用于判断大label适不适合放在聚类标签的那个概念下载, 具体最合适的是哪个概念
"""


import json
import copy
import argparse
import pandas as pd
from multiprocessing import Pool

# ---------------------大模型体系相关模块---------------------------------------
from zhipuai import ZhipuAI
# ChatOpenAI 是专门设计用于与 ChatGPT 或其他 OpenAI 模型进行交互的
from langchain_openai import ChatOpenAI
from langchain.prompts import (
    # 这是一个通用的聊天提示模板,用于定义整个聊天交互的结构
    # 它可以包含系统消息、用户消息、AI消息以及任何其他相关的元素
    ChatPromptTemplate,
    # MessagesPlaceholder可以作为一个占位符,在ChatPromptTemplate或其他模板中使用
    # 以便在运行时插入动态的消息列表,这增加了模板的灵活性和可定制性
    MessagesPlaceholder,
    # 专门用于定义系统级别的消息,这些消息通常不是由用户或AI直接发出的
    # 而是由系统生成,用于提供额外的上下文或指导
    SystemMessagePromptTemplate,
    # 人工消息提示模板,这是用户发送的消息
    HumanMessagePromptTemplate,
    # AI 消息提示模板,这是从 AI 发送的消息
    # 用于定义AI（通常是聊天模型）在对话中发送的消息格式
    AIMessagePromptTemplate,
)
# OutputParser 类解析 LLM 调用的输出:
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.output_parsers import StrOutputParser



############################### Prompt Template ################################

def prompt_generate(system, ai):
    """
        生成链中需要的Prompt模板
    """
    prompt = ChatPromptTemplate(
        messages=[
            # 定义一个系统消息
            SystemMessagePromptTemplate.from_template(system),
            # 定义一个AI消息
            AIMessagePromptTemplate.from_template(ai),
            # 定义一个人类消息
            HumanMessagePromptTemplate.from_template("{topic}")
        ]
    )

    return prompt


################################### Chains #####################################
"""
    模型温度:
        低温度(例如0.1):模型可能会生成非常标准、常见的回答,缺乏新意
        中温度(例如0.5-0.9):
            模型通常能在准确性和创造性之间找到一个良好的平衡,适用于大多数情况
        高温度(例如0.95或更高):
            模型可能会生成更加独特、有趣的回答,但也可能包含更多的错误或不连贯之处
"""


def chat_glm(prompt, user, parser):
    """
        输入提示词,运行glm-4-flash模型的一个链,然后格式化输出
    """
    # 创建一个ChatOpenAI实例
    llm = ChatOpenAI(
        # 设置模型的"温度"参数，影响生成内容的随机性和创造性
        temperature=0.5,
        # 设置使用的模型
        model="glm-4-flash",
        # 这里API Key
        openai_api_key="e7794ecfd23516cd297dc83e7f57a0e9.1ta3BPaDit6LuXJj",
        # 这里API Base
        openai_api_base="https://open.bigmodel.cn/api/paas/v4/"
    )
    # 定义个langchain的链
    chain = prompt | llm | parser
    # 调用链，生成链的结果
    response = chain.invoke({"topic": user})

    return response


##################################### Main #####################################
def label_to_concepts_tree_chatglm4(prompt, user_prompt, jsonparser, label):

    break_number = 1
    input_summary = ""
    while True:  # 如果结果报错，那么就重新用大模型生成，直到不报错为止
        #if break_number == 1000:  #增加重新运行的次数
        #    break
        try:
            input_summary = chat_glm(
                prompt, user_prompt, jsonparser
            )
            print("\n\n" + label + "\n大模型的结果是:\n\t",input_summary)
            break
        except Exception as e:
            # print(e)
            break_number = break_number + 1
            pass

    return [label, input_summary]


def multi_exec(input_file, output_path, multi):
    with open(input_file) as file:
        inpupt_dict = json.load(file)
    # 创建结果解析器
    jsonparser = StrOutputParser()
    system = [
        "你现在是一个生物医学专家，用户输入的有 功能标签 和 概念树中的概念，",
        "请判断功能标签适不适合归类到概念树的某一个概念下，直接归类或间接归类都可以，如果适合，给出最适合的的概念，",
        "输出结果需要按标准的Json格式输出。"
    ]
    ai = [
        "请输出标准的Json格式。\n",
        "如果适合放在概念树中，请输出： ",
        """
            is_suitable: True,
            concept: 这里填最适合的概念
        """,
        "如果不适合放在概念树的某个概念中，请输出： ",
        """
            is_suitable: False,
            concept: 这里填空字符串
            reason: 这里填原因
        """
    ]
    tasks = []
    for label,concepts in inpupt_dict.items():
        concepts_copy = copy.deepcopy(concepts)
        concepts_copy.remove(label)

        user_prompt = (
            "\n\n功能标签是：" + label + "\n\n",
            "\n\n概念树中的概念列表是: " + str(concepts_copy) + "\n\n"
        )
        prompt = prompt_generate(
            " ".join(system),
            " ".join(ai))
        task = (prompt, user_prompt, jsonparser, label)
        tasks.append(task)

    with Pool(multi) as pool:
        result = pool.starmap(label_to_concepts_tree_chatglm4, tasks)

    result_dict = {}
    for i in result:
        result_dict[i[0]] = i[1]

    with open(output_path, mode="w", encoding="utf-8") as file:
        json.dump(result_dict, file, indent=4)


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("-input","--input_path")
    parser.add_argument("-output","--output_path")
    parser.add_argument("-multi","--multi")
    args = parser.parse_args()

    return args


if __name__ == "__main__":
    params = parse_args()
    input_file = params.input_path
    output_path = params.output_path
    multi = int(params.multi)
    # input_file = "label_cluster_concept_tree_1.json"
    # output_path ="07.label_to_tree.json"
    multi_exec(input_file, output_path, multi)

